package com.whiteseason.spark.streaming

import org.apache.kafka.clients.consumer.ConsumerRecord
import org.apache.kafka.common.TopicPartition
import org.apache.kafka.common.serialization.StringDeserializer
import org.apache.spark.streaming.dstream.InputDStream
import org.apache.spark.streaming.kafka010._
import org.apache.spark.streaming.{Seconds, StreamingContext}
import org.apache.spark.{SparkConf, SparkContext}
import redis.clients.jedis.{Jedis, JedisPool, JedisPoolConfig, Transaction}

import scala.collection.JavaConverters._
import scala.collection.mutable
import scala.util.Try

object KafkaSparkStreamingDemo {

  def getOffset(topic: String, groupId: String): mutable.Map[TopicPartition, Long] = {
    val fromOffset = scala.collection.mutable.Map[TopicPartition, Long]()
    val jedis: Jedis = getConnection
    val keys = jedis.keys(s"offset_${groupId}_${topic}*")
    if (! keys.isEmpty) {
      keys.asScala.foreach(key => {
        val offset = jedis.get(key)
        val partition = Try(key.split(s"offset_${groupId}_${topic}_").apply(1)).getOrElse("0")
        fromOffset.put(new TopicPartition(topic, partition.toInt), offset.toLong)
      })
    }
    jedis.close()
    fromOffset
  }

  def main(args: Array[String]): Unit = {
    val conf = new SparkConf().setMaster("local[2]").setAppName("kafkaSparkStreaming")
    val context = new SparkContext(conf)
    context.setLogLevel("WARN")
    //获取上下文对象
    val ssc = new StreamingContext(context , Seconds(5))

    var locationStrategy: LocationStrategy = LocationStrategies.PreferConsistent
    val brokers = "spark:9092"
    val topic = "spark-kafka"
    val group = "sparkaGroup"
    val offsets = getOffset(topic,group)
    val kafkaParam = Map(
    "bootstrap.servers"-> brokers,
    "key.deserializer" ->classOf[StringDeserializer],
    "value.deserializer"->classOf[StringDeserializer],
    "group.id"->group,
    "auto.offset.reset"-> "latest",
    "enable.auto.commit" ->(false:java.lang.Boolean)
    );


    var consumerStrategy: ConsumerStrategy[String, String] = ConsumerStrategies.Subscribe(Array(topic), kafkaParam, offsets)

    var resultDStream: InputDStream[ConsumerRecord[String, String]] = KafkaUtils.createDirectStream(ssc, locationStrategy, consumerStrategy)

    resultDStream.foreachRDD(iter=>{
      if(iter.count() > 0){
        iter.foreach(record =>{
          val value : String = record.value()
          //record.offset()
          println(value)
        })
        val offsetRanges: Array[OffsetRange] = iter.asInstanceOf[HasOffsetRanges].offsetRanges
        if (!iter.isEmpty()) {
          val jedis: Jedis = getConnection
          val transaction: Transaction = jedis.multi()
          offsetRanges.foreach(
            iter=>{
              val key = s"offset_${group}_${iter.topic}_${iter.partition}"
              val value = iter.untilOffset
              transaction.set(key,value.toString)
            }
          )
          transaction.exec()
          transaction.close()
          jedis.close()
        }
      }

    })

    ssc.start()
    ssc.awaitTermination()

  }
  def getConnection: Jedis ={
    val redisconf = new JedisPoolConfig()
    val pool = new JedisPool(redisconf,"spark",6379)
    redisconf.setMaxTotal(20)
    redisconf.setMaxIdle(20)
    val jedis = pool.getResource
    jedis.auth("whiteseason")
    jedis
  }
}
